import numpy as np

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.io import Dataset, BatchSampler, DataLoader

# from paddle.fluid.dataloader import default_collate_fn

BATCH_NUM = 20
BATCH_SIZE = 32
EPOCH_NUM = 4

IMAGE_SIZE = 784
CLASS_NUM = 10

USE_GPU = False # whether use GPU to run model

bat = None

# define a random dataset
class RandomDataset(Dataset):
    def __init__(self, num_samples):
        self.num_samples = num_samples

    def __getitem__(self, idx):
        image = np.random.random([IMAGE_SIZE]).astype('float32')
        label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
        return image, label

    def __len__(self):
        return self.num_samples

dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)

cnt = 0
def collate_fn(batch):
    # image, labels = zip(*batch)
    # image = paddle.stack(images, 0)
    # labels = paddle.stack(labels, 0)
    # return batch
    return (
        paddle.zeros(shape=(32, 784)),
        paddle.zeros(shape=(32, 1)),
        paddle.ones(shape=(32, 42))
    )

class SimpleNet(nn.Layer):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(IMAGE_SIZE, CLASS_NUM)

    def forward(self, image, label=None):
        return self.fc(image)

simple_net = SimpleNet()
opt = paddle.optimizer.SGD(learning_rate=1e-3,
                          parameters=simple_net.parameters())
loader = DataLoader(dataset,
                    batch_size=BATCH_SIZE,
                    shuffle=True,
                    drop_last=True,
                    num_workers=2,
                    collate_fn=collate_fn,
)

for e in range(EPOCH_NUM):
    for i, (image, label, third) in enumerate(loader()):
        print(image.numpy().shape, label.numpy().shape)
        print(image.numpy())
        print(third)
